from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
import numpy as np
from ultralytics import YOLO

# 加载YOLOv8模型（这里使用预训练模型，也可以用自己训练的）
model = YOLO('yolov8n.pt')  # 这是一个轻量级模型，适合初学者

# 创建FastAPI应用
app = FastAPI(title="水果识别系统")

# 配置跨域，解决前后端通信问题
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源访问，实际应用中可以指定具体域名
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法
    allow_headers=["*"],  # 允许所有HTTP头
)

# 定义识别水果的接口
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # 读取上传的图片
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    # 转换为RGB格式（确保图片格式正确）
    if image.mode != "RGB":
        image = image.convert("RGB")
    
    # 使用YOLOv8模型进行预测
    results = model(image)
    
    # 处理识别结果
    predictions = []
    for result in results:
        for box in result.boxes:
            # 获取类别名称
            class_name = model.names[int(box.cls)]
            # 只保留水果类别（YOLOv8预训练模型中的水果类别）
            fruits = ['apple', 'banana', 'orange', 'grape', 'watermelon']
            if class_name in fruits:
                predictions.append({
                    "fruit": class_name,
                    "confidence": float(box.conf),  # 置信度，0-1之间，越接近1越准确
                    "box": box.xyxy.tolist()[0]  # 识别框的坐标
                })
    print(predictions)
    return {"predictions": predictions}

# 根路径，用于测试服务是否正常运行
@app.get("/")
async def root():
    return {"message": "水果识别系统正在运行中！请访问/predict接口进行识别。"}
